﻿#include "yolobase/yolo.h"
#include "yolo8/yolov8.h"
#include "yolobase/utils.h"

void setParameters(yolo::YoloParams& param)
{
	param.class_names = {"UUV"};
	param.num_class = 1; // for uuv
	param.model_file = "data/yolov8nsim_int32_8616.trt";

	param.batch_size = 1; //不能太大，显存溢出

	param.dst_w = 640;
	param.dst_h = 640;
	param.input_output_names = { "images",  "output0" };
	param.conf_thresh = 0.25f;
	param.iou_thresh = 0.45f;
}

int main(int argc, char** argv)
{
	/************************************************************************************************
	* init
	*************************************************************************************************/
	std::string video_path = "data/org3.mp4";

	// parameters
	yolo::YoloParams param;
	setParameters(param);

	bool is_show = true;

	cv::VideoCapture capture;
	capture.open(video_path);

	param.src_h = (int)capture.get(cv::CAP_PROP_FRAME_HEIGHT);
	param.src_w = (int)capture.get(cv::CAP_PROP_FRAME_WIDTH);

	yolo::YOLOV8 yolo(param);

	assert(yolo.init(param.model_file));
	/************************************************************************************************
	* recycle
	*************************************************************************************************/
	std::vector<cv::Mat> imgs_batch;
	imgs_batch.reserve(param.batch_size);

	cv::Mat frame;
	while (capture.isOpened())
	{
		if (imgs_batch.size() < param.batch_size) // get input
		{
			capture.read(frame);

			if (frame.empty())
			{
				std::cout << "no more video or camera frame" << std::endl;
				imgs_batch.clear(); // clear
				break;
			}
			else
			{
				imgs_batch.emplace_back(frame.clone());
			}
		}
		else // infer
		{
			auto found = yolo.detect(imgs_batch);
			if (is_show)
				utils::show(found, param.class_names, 1, imgs_batch);

			imgs_batch.clear(); // clear
		}
	}

	return  0;
}

